import os
import openai
from openai import *
import json
import argparse
from enum import Enum
from subprocess import Popen,PIPE
from openai.types.chat import *
from openai.types.chat.chat_completion_message_tool_call import Function
from openai.types.chat.chat_completion_chunk import *

class LogLevel(Enum):
    """日志级别枚举类，定义不同级别的日志类型和比较方法"""
    Debug = 1
    Info = 2
    Warn = 3
    Error = 4
    Fatal = 5
    def __eq__(self, value):
        if not isinstance(value, LogLevel):
            return False
        return self.value == value.value
        
    def __lt__(self, value):
        if not isinstance(value, LogLevel):
            return False
        return self.value < value.value
        
    def __le__(self, value):
        if not isinstance(value, LogLevel):
            return False
        return self.value <= value.value
        
    def __gt__(self, value):
        if not isinstance(value, LogLevel):
            return False
        return self.value > value.value
        
    def __ge__(self, value):
        if not isinstance(value, LogLevel):
            return False
        return self.value >= value.value
    
class Plugin:
    """Agent可调用的任何工具的抽象，管理各工具的提示词、定义语句、调用等功能"""
    def __init__(self, dir: str, name: str, timeout:float|None=None):
        """初始化插件实例
        
        Args:
            dir: 插件所在目录
            name: 插件文件名 
            timeout: 超时时间(秒)，None表示无超时
        """
        self.name = name
        self.dir = dir
        self.timeout = timeout

        self.refresh()
    
    def __call__(self, arguments: str, log_level: LogLevel=LogLevel.Info):
        """调用插件执行
        
        Args:
            arguments: 调用参数字符串
            log_level: 日志级别, 默认Info
            
        Returns:
            str: 包含stdout、stderr和返回码的格式化字符串
        """
        if log_level <= LogLevel.Info:
            print(f"+ Calling plugin: {self.name}({arguments})")
        #todo: 支持用户与插件交互
        proc = Popen(
            [self.name, arguments],
            executable=f"{self.dir}/{self.name}", 
            stderr=PIPE,
            )

        return_code = proc.wait(self.timeout)

        _,errinfo = proc.communicate()
        output = errinfo.decode().strip()

        if return_code != 0 and log_level <= LogLevel.Warn:
            for line in output.splitlines():
                print(f"+ERR+ {line}")
            print(f"+return_code+ {return_code}")
        elif log_level <= LogLevel.Debug:
            for line in output.splitlines():
                print(f"+ {line}")

        if return_code != 0:
            output = f"<error>{output}</error><return_code>{return_code}</return_code>"
        return output
    
    def refresh(self):
        """刷新插件信息，包括帮助信息和JSON定义"""
        proc_system = Popen(
            [self.name, "--help"],
            executable=f"{self.dir}/{self.name}", 
            stdout=PIPE,
            stderr=PIPE,
            )

        proc_definition = Popen(
            [self.name, "--usage"],
            executable=f"{self.dir}/{self.name}", 
            stdout=PIPE,
            stderr=PIPE,
            )

        if proc_system.wait(self.timeout) != 0:
            _,errinfo = proc_system.communicate()
            proc_system.kill()
            raise ChildProcessError(errinfo.decode())

        if proc_definition.wait(self.timeout) != 0:
            _,errinfo = proc_definition.communicate()
            proc_definition.kill()
            raise ChildProcessError(errinfo.decode())
        
        outinfo,errinfo = proc_system.communicate()
        self.system_prompt = outinfo.decode()

        outinfo,errinfo = proc_definition.communicate()
        self.definition = json.loads(outinfo.decode())

class PluginPool:
    """插件池管理类，负责加载和管理所有可用插件"""
    def __init__(self, tools_dir: str):
        """初始化插件池
        
        Args:
            tools_dir: 工具目录路径
        """
        self.tools_dir = tools_dir
        self.plugins: dict[str,Plugin] = {}
        self.reload_plugins()
        
    def reload_plugins(self, log_level=LogLevel.Info):
        """重新加载工具目录中的所有插件"""
        plugins = {}
        for filename in os.listdir(self.tools_dir):
            path = f"{self.tools_dir}/{filename}"
            if not os.path.isfile(path):
                continue
            if not os.access(path, os.X_OK):
                continue
            if self.plugins.__contains__(filename):
                plugins[filename] = self.plugins[filename]
            else:
                plugins[filename] = Plugin(self.tools_dir, filename)
        self.plugins = plugins
        if log_level <= LogLevel.Info:
            print(" * plugins reloaded")
        
    def __getitem__(self, name: str):
        return self.plugins.get(name)
    
    def __len__(self):
        return len(self.plugins)
    
    def __iter__(self):
        return self.plugins.__iter__()

class MiniAgent:
    """AI代理主类，负责核心逻辑和交互流程"""
    def __init__(self, config_path, tools_dir):
        self.running = True
        self.short_term_memory:list[ChatCompletionMessageParam] = []
        self.config = self.load_config(config_path)
        self.setup_model()
        self.plugins = PluginPool(tools_dir)

    def load_config(self, config_path):
        """加载并验证配置文件"""
        try:
            with open(config_path, 'r') as f:
                config = json.load(f)
            
            if not all(key in config for key in ['main_model', 'providers']):
                raise ValueError("配置文件缺少必要字段")
                
            main_model_config = config['main_model']
            provider_name = main_model_config['provider']
            
            if provider_name not in config['providers']:
                raise ValueError(f"未找到{provider_name}的提供商配置")
                
            return config
            
        except Exception as e:
            print(f"配置文件加载失败: {str(e)}")
            os._exit(1)

    def setup_model(self):
        """根据配置初始化模型
        
        从配置文件中读取模型提供商信息并初始化OpenAI客户端
        """
        provider_name = self.config['main_model']['provider']
        provider_config = self.config['providers'][provider_name]
        self.openai = openai.OpenAI(
            api_key=provider_config['api_key'],
            base_url=provider_config['base_url'],
        )

    def ask_model(self, user_input:str=None):
        """使用配置的模型处理用户输入，并更新短期记忆
        
        Args:
            user_input: 用户输入文本
            
        处理模型响应并维护对话上下文：
        1. 普通响应直接输出
        2. 流式响应实时输出
        """
        response = None
        try:
            self.plugins.reload_plugins()
            # 从配置中获取参数并设置stream默认值
            args:dict = self.config['main_model'].copy()

            args.pop('provider')

            messages:list = self.config['main_model'].get('messages',[])
            messages.extend([{"role": "system", "content":self.plugins[name].system_prompt} for name in self.plugins])

            
            if user_input is not None:
                self.short_term_memory.append({"role": "user", "content": user_input})
            messages.extend(self.short_term_memory)
            args['messages'] = messages

            args['tools'] = [self.plugins[name].definition for name in self.plugins]
            # print(self.short_term_memory)
            response = self.openai.chat.completions.create(**args)

            if isinstance(response, ChatCompletion):
                # 普通响应处理
                print("Agent: "+response.choices[0].message.content.strip())
                self.short_term_memory.append(response.choices[0].message.to_dict())
                return

            if isinstance(response, Stream):
                # 流式输出处理
                combined = ChatCompletionAssistantMessageParam()
                self.short_term_memory.append(combined)
                has_content = False
                for chunk in response:
                    chunk: ChatCompletionChunk
                    
                    for k,v in chunk.choices[0].delta:
                        if v is None:
                            continue
                        
                        if k == 'tool_calls':
                            if combined.get('tool_calls') is None:
                                combined['tool_calls'] = []
                            # combined['tool_calls'][v.]
                            for tool_call in v:
                                if not isinstance(tool_call, ChoiceDeltaToolCall):
                                    raise ValueError(f"unknown tool call: {tool_call}")
                                if not isinstance(tool_call.function, ChoiceDeltaToolCallFunction):
                                    raise ValueError(f"unknown tool call function: {tool_call.function}")
                                
                                if len(combined['tool_calls']) <= tool_call.index:
                                    combined['tool_calls'].extend([None for _ in range(tool_call.index - len(combined['tool_calls']) + 1)])
                                if combined['tool_calls'][tool_call.index] is None:
                                    combined['tool_calls'][tool_call.index] = ChatCompletionMessageToolCall(id="",function=Function(arguments="",name=""),type='function')

                                message: ChatCompletionMessageToolCall = combined['tool_calls'][tool_call.index]
                                
                                if tool_call.function.name is not None:
                                    if message.function.name is None:
                                        message.function.name = tool_call.function.name
                                    else:
                                        message.function.name += tool_call.function.name

                                if tool_call.function.arguments is not None:
                                    if message.function.arguments is None:
                                        message.function.arguments = tool_call.function.arguments
                                    else:
                                        message.function.arguments += tool_call.function.arguments

                                if tool_call.id is not None:
                                    if message.id is None:
                                        message.id = tool_call.id
                                    else:
                                        message.id += tool_call.id

                            continue
                        if combined.get(k) is None:
                            combined[k] = v
                        elif isinstance(v, str):
                            combined[k] += v
                        else:
                            raise ValueError(f'unexpected delta value: {k}:{v}')
                    if chunk.choices[0].delta.content:
                        content = chunk.choices[0].delta.content
                        if not has_content:
                            print("Agent: ",end='')
                        print(content, end="", flush=True)
                        has_content = True
                if has_content:
                    print()
                return

            raise TypeError(f"unexpected response type: {response}")
        except KeyboardInterrupt:
            if response is not None:
                response.close()
            print()
        except BaseException as e:
            print(args['messages'])
            print(f"API请求失败：")
            print(e)

    def call_plugin(self, arguments: str, name: str, id: str, use_function_call = False):
        """调用指定插件并处理结果
        
        Args:
            arguments: 调用参数
            name: 插件名称
            id: 工具调用ID
            use_function_call: 是否使用函数调用方式
            
        将插件调用结果添加到短期记忆中
        """
        plugin = self.plugins[name]
        if plugin is None:
            message = {
                "role": "tool",
                "content": "tool doesn't exist",
                "tool_call_id": id
            }
        else:
            message = {
                "role": "tool",
                "content": plugin(arguments),
                "tool_call_id": id
            }
        self.short_term_memory.append(message)

    def step(self, user_command: str):
        self.ask_model(user_command)
        try:
            while True:
                message = self.short_term_memory[-1]
                if message.__contains__('tool_calls'):
                    for tool_call in message['tool_calls']:
                        self.call_plugin(tool_call.function.arguments, tool_call.function.name, tool_call.id)
                    continue
                if message.__contains__('role') and message['role'] == 'tool':
                    self.ask_model()
                    continue
                break
        except KeyboardInterrupt:
            print()
            return
    def run(self):
        """启动agent的主循环"""
        print(f"{self.config['main_model']['model']} Agent 已启动（输入'EOF'结束）")
        while self.running:
            try:
                self.step(input(">>> "))
            except KeyboardInterrupt:
                print()
                continue
            except EOFError:
                print("\n检测到中断信号，正在退出...")
                self.running = False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='AI Agent')
    parser.add_argument('--config','-f', type=str, default="./config.json", required=False, help='配置文件路径 (例: ./config.json)')
    parser.add_argument('--tools','-t', type=str, default='./tools', required=False, help='工具文件夹路经 (例: ./tools)')
    parser.add_argument('--command','-c', type=str, required=False, help='让agent处理一条信息')
    args = parser.parse_args()
    agent = MiniAgent(args.config, args.tools)
    if args.command is None:
        agent.run()
    else:
        agent.step(args.command)
